/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024.
 * Description: RTOS Process Context asm header
 * Create: 2024-07-13
 */
#include <linux/bug.h>
#include <linux/sched.h>
#include <linux/rtos_rtpc.h>
#include <linux/rtos_rtpc_rpc.h>
#include <linux/preempt.h>
#include <linux/irqflags.h>

#include <asm/rtos_rtpc.h>
#ifdef CONFIG_ARM64
#include <asm/daifflags.h>
#endif

#include "rtpc.h"
#include "switch.h"

#ifdef CONFIG_RTOS_RTPC_RPC_DEBUG_BUILD
__attribute__((cold))
static void wake_up_check(struct task_struct *const tsk,
		struct rpc_service_info *const service_info)
{
	unsigned long flags;

	rtpc_spin_lock_irqsave(&service_info->lock, flags);
	BUG_ON(!service_info->need_exit);
	BUG_ON(service_info->calling_client_task);

	if (service_info->need_exit & (unsigned int)EXIT_FIXUP)
		BUG_ON(!service_info->server_need_wake ||
		       !service_info->calling_client_need_wake);

	if (tsk == service_info->server_task)
		BUG_ON(service_info->fixup_done);
	else
		BUG_ON(!service_info->fixup_done);

	rtpc_spin_unlock_irqrestore(&service_info->lock, flags);
}
#else
static __always_inline void wake_up_check(struct task_struct *t,
		struct rpc_service_info *s)
{
}
#endif

__attribute__((cold)) asmlinkage
long rtpc_wait_call_fast(struct task_struct *const tsk,
		struct rpc_service_info *const service_info)
{
	BUG_ON(tsk->audit_context && tsk->audit_context->in_syscall != 0);

	if (service_info->need_exit) {
		rtpc_spin_unlock(&service_info->lock);
		return -ENOLINK;
	}

	BUG_ON(service_info->server_init_state != SERVICE_CREATED);

	service_info->server_init_state = SERVER_START_SLEEP;

	rtpc_user_exit();
	preempt_disable();

	/*
	 * We must set TASK_RTPC_UNWAKEABLE first, then unlock
	 * Otherwise, there will be concurrency problems.
	 */
	set_current_state(TASK_RTPC_UNWAKEABLE);

	rtpc_spin_unlock(&service_info->lock);

	/* irq will be enable in __schedule() */
	rtpc_schedule_preempt_disabled();
	/* irq is enabled */
#ifdef CONFIG_ARM64
	local_daif_mask();
#else
	local_irq_disable();
#endif

	rpc_notice("%d/%s server wake up not sleep\n", tsk->pid, tsk->comm);

	BUG_ON(tsk != current);
	BUG_ON(rpc_get_service_info_self(tsk) != service_info);
	wake_up_check(tsk, service_info);

	preempt_enable_no_resched();
	rtpc_user_enter(true);

	return -ENOLINK;
}

__attribute__((cold))
long rtpc_wake_handle(void)
{
#ifdef CONFIG_ARM64
	local_daif_mask();
#else
	local_irq_disable();
#endif

	wake_up_check(current, rpc_get_service_info_self(current));

	preempt_enable_no_resched_notrace();
	rtpc_user_enter(false);

	rpc_notice("%d/%s be woken\n", current->pid, current->comm);
	return -1;
}

__attribute__((cold)) noinline
long ret_call_fast_need_exit_handler(void)
{
	struct rpc_service_info *const service_info = rpc_get_service_info_self(current);
	struct task_struct *const server_task = service_info->server_task;

	rpc_notice("ret call fast with need_exit!\n");

	if (service_info->server_need_wake) {
		rpc_notice("wake up server %d/%s in ret call fast\n",
				server_task->pid, server_task->comm);

		wake_up_state(server_task, TASK_RTPC_UNWAKEABLE);
	}
	return -ENOLINK;
}

__attribute__((cold)) noinline
void do_fixup_need_exit_handler(void)
{
	struct rpc_service_info *const service_info = rpc_get_service_info_self(current);

	rpc_notice("detect need exit in fixup\n");

	if (service_info->calling_client_need_wake) {
		struct task_struct *const client_task = service_info->calling_client_task;

		rpc_notice("wake up client %d/%s in fixup!\n",
				client_task->pid, client_task->comm);

		/*
		 * We must use write once here, because in the entry of fixup,
		 * we read it without service_info locked
		 */
		WRITE_ONCE(service_info->calling_client_task, NULL);
		wake_up_state(client_task, TASK_RTPC_UNWAKEABLE);
	}
}

#ifdef CONFIG_RTOS_RTPC_RPC_DEBUG_BUILD
/*
 * Rules:
 * 1. if is in need fixup state, KILL/EXIT_NORAML flags cannot be set
 * 2. if EXIT_NORMAL is set, client must be NULL
 * 3. if EXIT_FIXUP is set, and fixup is done, then client must be NULL (same as 4 + 5)
 * 4. if EXIT_FIXUP is set，server_need_wake and client_need_wake must be set
 * 5. if client_need_wake is set and fixup is done，then client must be NULL
 * 6. It's impossible that client != NULL and server == NULL
 * 7. A task which in need fixup state (client or server), it will not runs to kill or exit
 * 8. A task will never see itself is calling_client_task(except arm64 ret_call_fast/fixup)
 */

#ifdef CONFIG_ARM64
struct ret_x0_x1 {
	u64 x0;
	u64 x1;
};

typedef struct ret_x0_x1 rtpc_2regs_ret;

#define MAKE_2REGS_RET(x, y) ((struct ret_x0_x1){(x), (y)})
#else
typedef unsigned long long rtpc_2regs_ret;

#define MAKE_2REGS_RET(r0, r1) ((unsigned long long)(u32)(r0) | ((unsigned long long)(r1) << 32))
#endif

__attribute__((cold)) asmlinkage
rtpc_2regs_ret rtpc_need_exit_check_do_call(void *unused_var,
		const struct rpc_service_info *const service_info)
{
	(void)unused_var;

	/* rule 4 */
	if (service_info->need_exit & (unsigned int)EXIT_FIXUP)
		BUG_ON(!service_info->server_need_wake || !service_info->calling_client_need_wake);

	return MAKE_2REGS_RET(-ENOLINK, (uintptr_t)service_info);
}

__attribute__((cold)) asmlinkage
rtpc_2regs_ret rtpc_need_exit_check_ret_call_exception(void *unused_var,
		struct rpc_service_info *const service_info)
{
	(void)unused_var;

	BUG_ON(!service_info->fixup_done);

	/* rule 2 */
	BUG_ON(service_info->need_exit & (unsigned int)EXIT_NORMAL);

	/* rule 3 */
	BUG_ON(service_info->need_exit & (unsigned int)EXIT_FIXUP);

	/* rule 5 */
	BUG_ON(service_info->calling_client_need_wake);

	return MAKE_2REGS_RET(-ENOLINK, (uintptr_t)service_info);
}

__attribute__((cold)) asmlinkage
rtpc_2regs_ret rtpc_need_exit_check_other(void *unused_var,
		struct rpc_service_info *const service_info)
{
	(void)unused_var;

	/* rule 6 */
	BUG_ON(!service_info->server_task);

	/* rule 2 */
	BUG_ON(service_info->need_exit & (unsigned int)EXIT_NORMAL);

	if (!service_info->fixup_done) {
		/* Now is in need fixup state */

		/* rule 1 */
		BUG_ON(service_info->need_exit &
		       (~((unsigned int)KILL_FIXUP | (unsigned int)EXIT_FIXUP)));

		/* rule 4 */
		if (service_info->need_exit & (unsigned int)EXIT_FIXUP)
			BUG_ON(!service_info->server_need_wake ||
			       !service_info->calling_client_need_wake);
	} else {
		/* Now fixup is done */

		/* rule 3 */
		BUG_ON(service_info->need_exit & (unsigned int)EXIT_FIXUP);

		/* rule 5 */
		BUG_ON(service_info->calling_client_need_wake);
	}

	return MAKE_2REGS_RET(-ENOLINK, (uintptr_t)service_info);
}
#endif
